#include <iostream>
#include <fstream>
#include <vector>
#include <random>
#include "llvm/Pass.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/IR/Instructions.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/ValueMapper.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.h"
#include "llvm-c/Core.h"
#include "llvm/ADT/APInt.h"
#include "llvm/IR/GlobalVariable.h"
#include "llvm/IR/Type.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/Transforms/IPO/PassManagerBuilder.h"
using namespace llvm;
namespace
{
    class FRProtector : public ModulePass
    {
        public:
            static char ID;
            FRProtector() : ModulePass(ID){}
            bool runOnModule(Module &M) override;
        };
}

char FRProtector::ID = 0;
static RegisterPass<FRProtector> X("FRProtector","shuffle func name with old func name LLVM pass");
Pass *createFRProtectorPass() { return new FRProtector(); }

int GetTotalFuncNum(Module &module)
{
    Module &M = module;
    int cnt = 0;
    for (Module::iterator i=M.begin(); i!=M.end(); i++)
        cnt++;
    return cnt;
}

void shuffleFuncName(Module &module)
{
    Module &M = module;

    int total = GetTotalFuncNum(M);
    Function* func[total];
    Function* ori_func[total];
    int i = 0;
    int cnt = 0;
    int iterator_func = 0;
    for (Function &I : M)
    {
        
        func[iterator_func] = &I;
        if(func[iterator_func]->empty())
        {
            ori_func[cnt++] = &I;
            errs() << ori_func[cnt-1]->getName().str() <<"\n";
        }
        
        iterator_func++;
    }
    
    std::random_device rd;
    std::mt19937 mt(rd());
    std::uniform_int_distribution<int> rand(0, iterator_func);
    int it2 = iterator_func;
    for(i=0; i<iterator_func-1; i++)
    {
        int randNum = rand(mt) %it2;
        Function* tmp = func[randNum];
        func[randNum] = func[it2-1];
        func[it2-1] = tmp;
        it2--;
    }
    
    i = iterator_func;
    for (int ii=0; ii!=iterator_func; ii++)
    {
        i--;

        if(!func[i]->empty())
        {
            errs()<<"Found a non-empty function:" << func[i]->getName()<<"\n";
            if (func[i]->getName() == "main"){continue;}
            std::string newFakeName="_"+ori_func[rand(mt)%(cnt-1)]->getName().str();
            errs() << "Substitute with new fake name: " << newFakeName << "\n";

            ValueToValueMapTy VMap;
            Function *new_F=llvm::CloneFunction(func[i],VMap,NULL);
            new_F->setName(newFakeName);
            func[i]->replaceAllUsesWith(new_F);
            func[i]->eraseFromParent();

        }
    }
}



bool FRProtector::runOnModule(Module &M)
{
    shuffleFuncName(M);
    
    return false;
}
